#include <cstring>
#include <string>
#include <cstdio>
#include <iostream>
#include <fstream>
#include <TFile.h>
#include <TTree.h>
#include <vector>
#include <TCanvas.h>
#include <TH1F.h>
#include <TLegend.h>
#include <TStyle.h>
#include <sstream>

/*
 * Script to compute discriminants for Signal and Background from .lhco files 
 */

void histogramCreationCombined(char** fileTree, char** fileHists);
float getPz(float eta, float pT);
float getDsignal(float Lsig, float Lbg);
float getDbackground(float Lsig, float Lbg);

int main(int argc, char **argv) {
  std::cout << "Program started!" << std::endl;
  histogramCreationCombined(&argv[1],&argv[3]);
  std::cout << "Program finished!" << std::endl;
}

void histogramCreationCombined(char** fileTree, char** fileHists) {
  
  std::cout << "Using histograms from file: " << fileHists[0] << " and " << fileHists[1] << std::endl;
  std::cout << "Calculating D for " << fileTree[0] << " and " << fileTree[1] << std::endl;

  /* Open file with data in "tree" */
  TFile *fTree = TFile::Open(fileTree[0],"READ");
  TTree **t;
  t = new TTree*[2];
  fTree->GetObject("tree",t[0]);
  TFile *fTree2 = TFile::Open(fileTree[1],"READ");
  fTree2->GetObject("tree",t[1]);
  int mass;
  /* TEMPLATES */
  /* Open file with histograms with signal and background histograms and load them */
  /* Background and 125 */
  TFile *fHists = TFile::Open(fileHists[0],"READ");
  /* 120 to 130 */
  int massPoints = 11;
  TFile *fHists2[massPoints];
  for (mass=0; mass<massPoints; mass++) {
    fHists2[mass] = TFile::Open(fileHists[1+mass],"READ");
  }
  const int NHISTSfile = 12;
  /* histsFile: histograms from file . 125 and background template */
  TH1F*** histsFile=new TH1F**[NHISTSfile];
  for (int ihist=0; ihist<NHISTSfile; ihist++) {
    histsFile[ihist]=new TH1F*[2];
  }
  /* Mass templates */
  TH1F**** histsFile2=new TH1F***[NHISTSfile];
  for (int ihist=0; ihist<NHISTSfile; ihist++) {
    histsFile2[ihist]=new TH1F**[2];
    for (mass=0; mass<massPoints; mass++) {
      histsFile2[ihist][0]=new TH1F*[massPoints];
      histsFile2[ihist][1]=new TH1F*[massPoints];
    }
  }
  
  const char * histStart = "h";
  std::ostringstream ss;
  for(int i = 0; i < NHISTSfile; i++) {
    for(int j = 0; j < 2; j++) {
      ss.str("");
      ss << histStart << i+1 << "_" << j;
      fHists->GetObject((ss.str()).c_str(), histsFile[i][j]);
    }
    for(int j = 0; j < 2; j++) {
      ss.str("");
      ss << histStart << i+1 << "_" << j;
      for (mass=0; mass<massPoints; mass++) {
	fHists2[mass]->GetObject((ss.str()).c_str(), histsFile2[i][j][mass]);
      }
    }
  }

  /* Style for histograms */
  TStyle* myStyle = new TStyle("Plain","Style1");
  TCanvas *c1 = new TCanvas("c1","MadGraph analysis",200,10,700,500);

  std::vector<int> *typ = 0;
  std::vector<float> *eta = 0;
  std::vector<float> *phi = 0;
  std::vector<float> *pt = 0;
  std::vector<float> *jmas = 0;
  std::vector<float> *ntrk = 0;
  std::vector<float> *btag = 0;
  std::vector<float> *hadEm = 0;

  TBranch *btyp;
  TBranch *beta;
  TBranch *bphi;
  TBranch *bjmas;
  TBranch *bntrk;
  TBranch *bpt;
  TBranch *bbtag;
  TBranch *bhadEm;

  /* Number of histograms: 0: signal, 1:background */
  int NHISTS = 2;
  /* New histograms for D */
  TH1F *hists[NHISTS];
  hists[0] = new TH1F("h1","Discriminant",50,0.0,1);
  hists[1] = new TH1F("h2","Discriminant",50,0.0,1);


  /* Weight: assumed ratio signal vs background */
  double r = 1;
  int weight_count = 0;
  int event_count = 0;
  int weight_count2 = 0;
  int event_count2 = 0;
  int rPoints=1;
  double weight[massPoints][10];
  
  /* Likelihood 0:signal 1:background*/
  // For Discriminant (from 125 template set)
  float L[2], D[2];
  /* Template likelihoods for all masses and background */
  float Lw[2][massPoints];
  

  //Impose a cut: 0 no cut, 1 cut
  int cut =0;
  double fractions[rPoints];
  for (int l= 0;l<rPoints; l++) {
    fractions[l]=1;
    //fractions[l]=0.562;
  }
  int nLeptons, pos = 0;
  float bosonMass, EFirstLepton, currentE,etaFirstLepton, currentEta, phiFirstLepton, currentPhi, dEta, dPhi = 0;
  float pTx, pTy, pz,p,pTFirstLepton, currentpT = 0;

 
  int eventsMax=0;
  const float muonMass = 0.105658;
 
  for (int rPoint= 0;rPoint<rPoints; rPoint++) {
  
    for (mass= 0;mass<massPoints; mass++) {
      weight[mass][rPoint] = 0;
      weight_count = 0;
      event_count = 0;
      weight_count2 = 0;
      event_count2 = 0;
      // std::cout << 120+mass<< " ";

    
      for(int nFiles=0; nFiles < 2; nFiles++) {
      
	// Signal
	if(nFiles==0) {
	  eventsMax = 5000;
	}
	// Background
	else if(nFiles==1) {
	  eventsMax = 5000;
	}

    
	t[nFiles]->SetBranchAddress("typ",&typ,&btyp);
	t[nFiles]->SetBranchAddress("eta",&eta,&beta);
	t[nFiles]->SetBranchAddress("phi",&phi,&bphi);
	t[nFiles]->SetBranchAddress("jmas",&jmas,&bjmas);
	t[nFiles]->SetBranchAddress("ntrk",&ntrk,&bntrk);
	t[nFiles]->SetBranchAddress("pt",&pt,&bpt);
	t[nFiles]->SetBranchAddress("btag",&btag,&bbtag);
	t[nFiles]->SetBranchAddress("hadEm",&hadEm,&bhadEm);
	int entriest = (int)t[nFiles]->GetEntries();

     

   
     
	for (int i = 0; i < eventsMax; i++) {
	  //entriest
	  long long tentry = t[nFiles]->LoadTree(i);
	  btyp->GetEntry(tentry);
	  beta->GetEntry(tentry);
	  bphi->GetEntry(tentry);
	  bpt->GetEntry(tentry);
	  bjmas->GetEntry(tentry);
	  bntrk->GetEntry(tentry);
	  bbtag->GetEntry(tentry);
	  bhadEm->GetEntry(tentry);
	  bosonMass = 0.0;
	  nLeptons = 0;
	  etaFirstLepton = 0.0;
	  currentEta = 0.0;
	  pTFirstLepton = 0.0;
	  currentpT = 0.0;
	  EFirstLepton = 0.0;
	  currentE = 0.0;
	  phiFirstLepton = 0.0;
	  currentPhi = 0.0;
	  D[0] = 1.0;
	  D[1] = 1.0;
	  L[0] = 1.0;
	  L[1] = 1.0;
	  Lw[0][mass] = 1.0;
	  Lw[1][mass] = 1.0;
	  cut=0;
    
	  for (unsigned int j = 0; j < typ->size(); ++j) {
	    if((typ->at(j))==2) {
	      //hists[5][k+NFILESLHE]->Fill(pt->at(j));
	      //hists[0][k+NFILESLHE]->Fill(eta->at(j));
	      //hists[2][k+NFILESLHE]->Fill(phi->at(j));
	      bosonMass += pt->at(j);
	      nLeptons++;
	      etaFirstLepton = currentEta;
	      currentEta = eta->at(j);
	      pTFirstLepton = currentpT;
	      EFirstLepton = currentE;

	      pz = getPz(eta->at(j), pt->at(j));
	      p = std::sqrt(pz*pz+pt->at(j)*pt->at(j));
	      currentE = std::sqrt(p*p+muonMass*muonMass);

	      currentpT = pt->at(j);
	      phiFirstLepton = currentPhi;
	      currentPhi = phi->at(j);
	    }
	    //hists[4][k+NFILESLHE]->Fill(typ->at(j));
	  }
	  if(nLeptons==2) {	

	    pTx = std::cos(phiFirstLepton)*pTFirstLepton+std::cos(currentPhi)*currentpT;
	    pTy = std::sin(phiFirstLepton)*pTFirstLepton+std::sin(currentPhi)*currentpT;
	    pz = getPz(etaFirstLepton,pTFirstLepton)+ getPz(currentEta,currentpT);
	    dEta = std::abs(etaFirstLepton-currentEta);
	    dPhi = std::abs(phiFirstLepton-currentPhi);
	    p = std::sqrt(pTx*pTx+pTy*pTy+pz*pz);


	    if(currentpT>=pTFirstLepton) {
	      if((currentpT<10||currentpT>130)||(pTFirstLepton<5||pTFirstLepton>70)) {
		cut=1;
		//std::cout << "Leading pT " << currentpT << std::endl;
	      }
	      for(int l=0;l<2;l++) {
		pos=histsFile[10][l]->FindBin(currentpT);
		L[l]*=histsFile[10][l]->GetBinContent(pos);
		pos=histsFile[11][l]->FindBin(pTFirstLepton);
		L[l]*=histsFile[11][l]->GetBinContent(pos);
		if(pos==101||pos==0) {
		  L[l]*=0.0;
		}
	      }
	      for(int l=0;l<2;l++) {
	 
		pos=histsFile2[10][l][mass]->FindBin(currentpT);
		Lw[l][mass]*=histsFile2[10][l][mass]->GetBinContent(pos);
		pos=histsFile2[11][l][mass]->FindBin(pTFirstLepton);
		Lw[l][mass]*=histsFile2[11][l][mass]->GetBinContent(pos);
		if(pos==101||pos==0) {
		  Lw[l][mass]*=0.0;
		}
	      }
	    }
	    else {
	      if((pTFirstLepton<10||pTFirstLepton>130)||(currentpT<5||currentpT>70)) {
		cut=1;
		//std::cout << "Leading pT " << pTFirstLepton << std::endl;
	      }
	      for(int l=0;l<2;l++) {
		pos=histsFile[10][l]->FindBin(pTFirstLepton);
		L[l]*=histsFile[10][l]->GetBinContent(pos);
		pos=histsFile[11][l]->FindBin(currentpT);
		L[l]*=histsFile[11][l]->GetBinContent(pos);
		if(pos==101||pos==0) {
		  L[l]*=0.0;
		}
	      }
	      for(int l=0;l<2;l++) {
		pos=histsFile2[10][l][mass]->FindBin(pTFirstLepton);
		Lw[l][mass]*=histsFile2[10][l][mass]->GetBinContent(pos);
		pos=histsFile2[11][l][mass]->FindBin(currentpT);
		Lw[l][mass]*=histsFile2[11][l][mass]->GetBinContent(pos);
		if(pos==101||pos==0) {
		  Lw[l][mass]*=0.0;
		}
	      }
	    }

	    r=fractions[rPoint];
	    D[0] =getDsignal(L[0], L[1]);
            if(nFiles==0) event_count++;
	    if(nFiles==1) event_count2++;
	    if(D[0]>=0.0&&Lw[0][mass]+Lw[1][mass]!=0.0&&nFiles==0&&cut==0) {
	      //std::cout << "Weight = " <<  -log(Lw[0]*r+Lw[1]*(1-r)) << std::endl;
	      weight[mass][rPoint] += -log(Lw[0][mass]*r+Lw[1][mass]*(1-r));
	      if(-log(Lw[0][mass]*r+Lw[1][mass]*(1-r))>1000) {
		std::cout << "Problem: " << i << std::endl;
	      }
	      weight_count++;
	    }
	    else if(D[0]>=0.0&&Lw[0][mass]+Lw[1][mass]!=0.0&&nFiles==1&&cut==0) {
	      weight[mass][rPoint] += -log(Lw[0][mass]*r+Lw[1][mass]*(1-r));
	      if(-log(Lw[0][mass]*r+Lw[1][mass]*(1-r))>1000) {
		std::cout << "Problem: " << i << std::endl;
	      }
	      weight_count2++;
	    }

	    hists[nFiles]->Fill(D[0]);
	    //hists[1]->Fill(D[1]);
     
	  }
	}
      }
      std::cout << weight[mass][rPoint] << " ";
 
    }
    std::cout << std::endl;
  }
  std::cout << "From " << weight_count << " events out of "<< event_count << std::endl;
  std::cout << "Bg: From " << weight_count2 << " events out of " << event_count2 << std::endl;


  for(int i = 0; i < NHISTS; i++) {
    hists[i]->Scale(1.0/hists[i]->Integral());
  }

  double maximum = 0.0;
  double lmax = 0.0;
  for(int i = 0; i < NHISTS; i++) {
    hists[i]->SetXTitle("D");
    hists[i]->SetYTitle("Fraction of events");
    lmax = hists[i]->GetMaximum();
    if(lmax>maximum) {
      maximum = lmax;
      }
  }

  const char* fileExtension = ".pdf";
  const char* fileStart = "outputCombined/hist";
  std::ostringstream oss;

  myStyle->cd();
  myStyle->SetOptStat(0);
  TLegend* leg = new TLegend(0.7,0.7,0.85,0.85);
  leg->SetFillColor(0);

  leg->SetHeader("");
  leg->AddEntry(hists[0],"H #rightarrow #mu #mu","l");
  leg->AddEntry(hists[1],"Z #rightarrow #mu #mu","l");

  for(int i = 0; i < NHISTS; i++) {
    hists[i]->SetMaximum(maximum*1.2);
  }
  hists[0]->SetLineColor(4);
  hists[1]->SetLineColor(2);
  hists[0]->Draw();
  hists[1]->Draw("same");
  leg->Draw();
  c1->Modified();
  c1->Update();
  oss.str("");
  oss << fileStart << 1 << fileExtension;
  c1->SaveAs((oss.str()).c_str(),"");

  t[0]->ResetBranchAddresses();
  t[1]->ResetBranchAddresses();
  

}

float getPz(float eta, float pT) {
  float theta = 2*std::atan(std::exp(-eta));
  if(eta>=0) {
    theta+=atan(0);
  }
  return pT/tan(theta);
}


float getDsignal(float Lsig, float Lbg) {
  float D;
  float R;
  if(Lbg==0.0) {
    D = 1.0;
  }
  else if(Lsig==0) {
    D = -1;
  }
  else {
    R = Lsig/Lbg;
    D = R/(R+1);
  }
  return D;
}

float getDbackground(float Lsig, float Lbg){
  float D;
  float R;
  if(Lsig==0.0) {
    D = 1.0;
  }
  else if(Lbg==0) {
    D = 0.5;
    std::cout << "Warning: Lsig and Lbg both zero, D=0.5!" << std::endl;
  }
  else {
    R = Lbg/Lsig;
    D = R/(R+1);
  }
  return D;
}
